#pragma once
#include <xpam/com/com.h>

namespace xpam_com
{
    namespace dshow
    {
        class sample_grabber_cb : public xpam_com::ISampleGrabberCB
        {
        protected:
            int _width, _height;
            xpam::dsp::frame _frame;
            double _time;
        public:
            sample_grabber_cb() : _width(0), _height(0)
            { }

            virtual ~sample_grabber_cb()
            { }

            void frame(xpam::dsp::frame & fr) { fr = _frame; }
            double time() { return _time; }

            // ISampleGrabberCB implementation
            STDMETHOD_(ULONG, AddRef)();
            STDMETHOD_(ULONG, Release)();
            STDMETHOD(QueryInterface)(REFIID riid, void ** ppv);
            HRESULT IUnknown::QueryInterface(const IID &,void **) {};
            STDMETHOD(raw_SampleCB)(double sampleTime, xpam_com::IMediaSample * sample);
            STDMETHOD(raw_BufferCB)(double sampleTime, unsigned char * buffer, long bufferSize);
        };

        STDMETHODIMP_(ULONG) sample_grabber_cb::AddRef() 
        { 
            return 0; 
        }

        STDMETHODIMP_(ULONG) sample_grabber_cb::Release() 
        { 
            return 0; 
        }

        STDMETHODIMP sample_grabber_cb::QueryInterface(REFIID riid, void ** ppv)
        {
            if (riid == xpam_com::IID_ISampleGrabberCB || riid == IID_IUnknown ) 
            {
                *ppv = (void *) static_cast<xpam_com::ISampleGrabberCB *>(this);
                return NOERROR;
            }    

            return E_NOINTERFACE;
        }

        STDMETHODIMP sample_grabber_cb::raw_SampleCB(double sampleTime, xpam_com::IMediaSample * sample)
        {
            unsigned char * buffer = NULL;
            sample->GetPointer(&buffer);
            _frame.assign(buffer, _width, _height, 3);
            _time = sampleTime;
            return S_OK;
        }

        STDMETHODIMP sample_grabber_cb::raw_BufferCB(double sampleTime, unsigned char * buffer, long bufferSize)
        {
            _frame.assign(buffer, _width, _height, 3);
            _time = sampleTime;
            return 0;
        }
    }
}